import json
import os
import tempfile

from flask import Flask, request, jsonify
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
import numpy as np

def predict_image(img_path, top_n=3):
    # 加载图像并进行预处理
    img = image.load_img(img_path, target_size=(150, 150))  # 根据您的模型输入大小调整
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = preprocess_input(img_array)  # 归一化处理

    # 进行预测
    predictions = model.predict(img_array)

    # 获取前 top_n 个预测结果及其置信度
    top_indices = np.argsort(predictions[0])[-top_n:][::-1]
    top_classes = [(class_labels[i], predictions[0][i]) for i in top_indices]

    return top_classes

def predict_top1(img_path):
    # 加载图像并进行预处理
    img = image.load_img(img_path, target_size=(150, 150))  # 根据您的模型输入大小调整
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = preprocess_input(img_array)  # 归一化处理

    # 进行预测
    predictions = model.predict(img_array)

    # 获取最高置信度的预测结果
    top_index = np.argmax(predictions[0])
    top_class = class_labels[top_index]
    confidence = predictions[0][top_index]

    return top_class, confidence

# 加载类别标签
def load_class_labels(file_path):
    with open(file_path, 'r', encoding='utf-8') as f:
        return json.load(f)

# 加载模型
model = load_model('model/mushroom_classification_model_20241222_233510.keras')  # 修改为您的模型文件名
class_labels = load_class_labels('model/class_labels.json')  # 确保该文件存在


app = Flask(__name__)

@app.route('/predict', methods=['POST'])
def predict():
    try:
        if 'image' not in request.files:
            return jsonify({'status': 'error', 'message': 'No image file provided'}), 400

        file = request.files['image']
        if file.filename == '':
            return jsonify({'status': 'error', 'message': 'No selected file'}), 400

        # 保存上传的图像文件到服务器的临时目录
        temp_dir = tempfile.gettempdir()
        file_path = os.path.join(temp_dir, file.filename)
        file.save(file_path)
        print('File saved to {}'.format(file_path))

        # 进行图像处理和预测
        predictions = predict_image(file_path, top_n=3)  # 返回多个预测结果
        print('Predicted labels: {}'.format(predictions))

        # 格式化预测结果
        results = [{'class': class_name, 'confidence': float(confidence)} for class_name, confidence in predictions]

        # 删除临时文件
        os.remove(file_path)

        return jsonify({'status': 'success', 'predictions': results})

    except Exception as e:
        return jsonify({'status': 'error', 'message': f'Error processing request: {str(e)}'}), 500

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=7000)
